package com.zhou.util;

import com.google.common.collect.Lists;
import com.zhou.po.CommunityChat;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.api.common.io.InputFormat;
import org.apache.flink.api.common.io.statistics.BaseStatistics;
import org.apache.flink.api.java.utils.ParameterTool;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.core.io.InputSplit;
import org.apache.flink.core.io.InputSplitAssigner;
import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.DeserializationFeature;
import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.util.Preconditions;
import org.apache.http.HttpHost;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.MatchQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.RangeQueryBuilder;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;

import java.io.IOException;
import java.net.UnknownHostException;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.List;
import java.util.Queue;
import java.util.concurrent.TimeUnit;

@Slf4j
public class ESInputFormat implements InputFormat<CommunityChat, RequestInputSplit> {

    private static final long serialVersionUID = -5952595992199493248L;

    private static RestHighLevelClient client=null;

    private int pageSize=100;

    BoolQueryBuilder boolQuery =null;

    private ObjectMapper mapper=null;

    public ESInputFormat(){
        mapper=new ObjectMapper();
        //mapper.setPropertyNamingStrategy(PropertyNamingStrategy.SNAKE_CASE);
        mapper.configure(DeserializationFeature.FAIL_ON_IGNORED_PROPERTIES,false);
        mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES,false);
    }


    public static RestHighLevelClient initClient(ParameterTool parameterTool) throws UnknownHostException {
        client = new RestHighLevelClient(
                RestClient.builder(new HttpHost(parameterTool.get(PropertiesKey.ES_IP), parameterTool.getInt(PropertiesKey.ES_PORT), "http")
                ));
        return client;
    }

    @Override
    public void configure(Configuration configuration) {

    }

    @Override
    public BaseStatistics getStatistics(BaseStatistics baseStatistics) throws IOException {
        return null;
    }

    @Override
    public RequestInputSplit[] createInputSplits(int i) throws IOException {
        SearchRequest searchRequest;
        searchRequest = new SearchRequest("community_im_chat");
        searchRequest.types("_doc");
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        sourceBuilder.size(0);

        sourceBuilder.timeout(new TimeValue(60, TimeUnit.SECONDS));
        MatchQueryBuilder chatTypeBuilder = QueryBuilders.matchQuery("chat_type", "chat");
        RangeQueryBuilder timestamp = QueryBuilders.rangeQuery("timestamp");
        timestamp.gte(1589817600000l);
        RangeQueryBuilder lt = timestamp.lt(1589904000000l);
        boolQuery = QueryBuilders.boolQuery();
        boolQuery.must().addAll(Lists.newArrayList(chatTypeBuilder, timestamp, QueryBuilders.existsQuery("chat_type")));
        sourceBuilder.query(boolQuery);
        searchRequest.source(sourceBuilder);
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        long totalHits = searchResponse.getHits().getTotalHits().value;
        log.info("es中查询的数据总量为{}",totalHits);
        int taotalNumber=(int)(totalHits-1)/pageSize+1;
        RequestInputSplit[] requests=new RequestInputSplit[taotalNumber];
        for (int pageNum=1;pageNum<=taotalNumber;pageNum++){
            SearchRequest request = new SearchRequest("community_im_chat");
            SearchSourceBuilder builder = new SearchSourceBuilder();
            builder.from((pageNum-1)*pageSize);
            builder.size(pageSize);
            builder.query(boolQuery);
            request.source(sourceBuilder);
            requests[pageNum-1]=new RequestInputSplit(pageNum-1);
        }
        return requests;
    }

    @Override
    public InputSplitAssigner getInputSplitAssigner(RequestInputSplit[] requestInputSplits) {
        return new RequestSplitAssigner(requestInputSplits);
    }

    private RequestInputSplit curInputSplit = null;

    private LinkedList<SearchHit> hits=null;


    @Override
    public void open(RequestInputSplit requestInputSplit) throws IOException {
        curInputSplit=requestInputSplit;
        SearchRequest request = new SearchRequest("community_im_chat");
        SearchSourceBuilder builder = new SearchSourceBuilder();
        builder.from(curInputSplit.getSplitNum()*pageSize);
        builder.size(pageSize);
        MatchQueryBuilder chatTypeBuilder = QueryBuilders.matchQuery("chat_type", "chat");
        RangeQueryBuilder timestamp = QueryBuilders.rangeQuery("timestamp");
        timestamp.gte(1589817600000l);
        RangeQueryBuilder lt = timestamp.lt(1589904000000l);
        boolQuery = QueryBuilders.boolQuery();
        boolQuery.must().addAll(Lists.newArrayList(chatTypeBuilder, timestamp, QueryBuilders.existsQuery("chat_type")));
        builder.query(boolQuery);
        request.source(builder);
        log.info(builder.toString());
        SearchResponse searchResponse = client.search(request, RequestOptions.DEFAULT);
        SearchHit[] searchHits = searchResponse.getHits().getHits();
        hits=new LinkedList(Arrays.asList(searchHits));
    }

    @Override
    public boolean reachedEnd() throws IOException {
        return curInputSplit == null;
    }

    @Override
    public CommunityChat nextRecord(CommunityChat s) throws IOException {
        SearchHit pop = hits.pop();
        if (hits.size()==0){
            curInputSplit = null;
        }
        String sourceAsString = pop.getSourceAsString();
        CommunityChat communityChat = mapper.readValue(sourceAsString, CommunityChat.class);
        return communityChat;
    }

    @Override
    public void close() throws IOException {

    }

    private class RequestSplitAssigner implements InputSplitAssigner{

        private Queue<RequestInputSplit> splits;

        public RequestSplitAssigner(RequestInputSplit[] requests){
            splits=new LinkedList<>(Arrays.asList(requests));
        }

        @Override
        public RequestInputSplit getNextInputSplit(String s, int i) {
            log.info("Getting copy task for task: " + i);
            return splits.poll();
        }

        @Override
        public void returnInputSplit(List<InputSplit> list, int i) {
            synchronized (this.splits) {
                for (InputSplit split : list) {
                    Preconditions.checkState(this.splits.add((RequestInputSplit)split));
                }
            }
        }
    }


    /*private static RestHighLevelClient client=null;

    private int pageSize=1000;


    public static RestHighLevelClient initClient(ParameterTool parameterTool) throws UnknownHostException {
        client = new RestHighLevelClient(
                RestClient.builder(new HttpHost(parameterTool.get(PropertiesKey.ES_IP), parameterTool.getInt(PropertiesKey.ES_PORT), "http")
                ));
        return client;
    }

    @Override
    public void configure(Configuration configuration) {

    }

    @Override
    public BaseStatistics getStatistics(BaseStatistics baseStatistics) throws IOException {
        return null;
    }

    @Override
    public RequestInputSplit[] createInputSplits(int i) throws IOException {
        RequestInputSplit[] requests=new RequestInputSplit[];
        SearchRequest searchRequest = new SearchRequest("community_im_chat");
        searchRequest.types("_doc");
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        sourceBuilder.size(0);
        MatchQueryBuilder chatTypeBuilder = QueryBuilders.matchQuery("chat_type", "chat");
        RangeQueryBuilder timestamp = QueryBuilders.rangeQuery("timestamp");
        timestamp.gte(1588694400000l);
        RangeQueryBuilder lt = timestamp.lt(1588780800000l);
        BoolQueryBuilder boolQuery = QueryBuilders.boolQuery();
        boolQuery.must().addAll(Lists.newArrayList(chatTypeBuilder, timestamp, QueryBuilders.existsQuery("chat_type")));
        sourceBuilder.timeout(new TimeValue(60, TimeUnit.SECONDS));
        sourceBuilder.query(boolQuery);
        searchRequest.source(sourceBuilder);
        SearchResponse searchResponse = client.search(searchRequest);
        long totalHits = searchResponse.getHits().getTotalHits();
        int taotalNumber=(int)(totalHits-1)/pageSize+1;

        for (int pageNum=1;i<=taotalNumber;pageNum++){
            SearchRequest request = new SearchRequest("community_im_chat");
            SearchSourceBuilder builder = new SearchSourceBuilder();
            builder.from((pageNum-1)*pageSize);
            builder.size(pageSize);
            builder.query(boolQuery);
            request.source(sourceBuilder);
            requests[pageNum-1]=new RequestInputSplit(request,pageNum-1);
        }
        return requests;
    }

    @Override
    public InputSplitAssigner getInputSplitAssigner(FileCopyTaskInputSplit[] fileCopyTaskInputSplits) {
        return null;
    }

    @Override
    public void open(FileCopyTaskInputSplit fileCopyTaskInputSplit) throws IOException {

    }

    @Override
    public boolean reachedEnd() throws IOException {
        return false;
    }

    @Override
    public FileCopyTask nextRecord(FileCopyTask fileCopyTask) throws IOException {
        return null;
    }

    @Override
    public void close() throws IOException {

    }


    private class FileCopyTaskAssigner implements InputSplitAssigner {
        private Queue<FileCopyTaskInputSplit> splits;

        public FileCopyTaskAssigner(FileCopyTaskInputSplit[] inputSplits) {
            splits = new LinkedList<>(Arrays.asList(inputSplits));
        }

        @Override
        public InputSplit getNextInputSplit(String host, int taskId) {
            LOGGER.info("Getting copy task for task: " + taskId);
            return splits.poll();
        }

        @Override
        public void returnInputSplit(List<InputSplit> splits, int taskId) {
            synchronized (this.splits) {
                for (InputSplit split : splits) {
                    Preconditions.checkState(this.splits.add((FileCopyTaskInputSplit) split));
                }
            }
        }
    }*/
}
